package io.kiki.sba.registry.server.metadata.slot.assigner;


import io.kiki.sba.registry.common.model.slot.DataStoreNode;
import io.kiki.sba.registry.common.model.slot.Slots;
import io.kiki.sba.registry.exception.SofaRegistryRuntimeException;
import io.kiki.sba.registry.server.metadata.monitor.Metrics;
import io.kiki.sba.registry.server.metadata.slot.SlotAssigner;
import io.kiki.sba.registry.server.metadata.slot.balance.BalancePolicy;
import io.kiki.sba.registry.server.metadata.slot.balance.BalancePolicyImpl;
import io.kiki.sba.registry.server.metadata.slot.util.FollowerToAssign;
import io.kiki.sba.registry.server.metadata.slot.util.MigratableSlotInfo;
import io.kiki.sba.registry.server.metadata.slot.util.builder.SlotsBuilder;
import io.kiki.sba.registry.server.metadata.slot.util.comparator.Comparators;
import io.kiki.sba.registry.server.metadata.slot.util.selector.Selectors;
import io.kiki.sba.registry.util.MathUtils;
import lombok.Getter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;

@Getter
public class DefaultSlotAssigner implements SlotAssigner {

    private static final Logger logger = LoggerFactory.getLogger(DefaultSlotAssigner.class);
    protected final SlotsBuilder slotsBuilder;
    private final Set<String> dataStoreNodeIdSet;
    private final MigratableSlotInfo migratableSlotInfo;

    public DefaultSlotAssigner(SlotsBuilder slotsBuilder, Collection<String> dataStoreNodeIdSet) {
        this.slotsBuilder = slotsBuilder;
        this.dataStoreNodeIdSet = Collections.unmodifiableSet(new TreeSet<>(dataStoreNodeIdSet));
        this.migratableSlotInfo = slotsBuilder.getAssignableSlots();
    }

    @Override
    public Slots assign() {
        if (migratableSlotInfo.isEmpty() || dataStoreNodeIdSet.isEmpty()) {
            throw new SofaRegistryRuntimeException("no pending slot or available dataServers for reassignment");
        }
        BalancePolicy balancePolicy = new BalancePolicyImpl();
        final int ceilAvg = MathUtils.divideCeil(slotsBuilder.getSlotNums(), dataStoreNodeIdSet.size());
        final int highWaterMarkSlotLeaderNums = balancePolicy.getHighWaterMarkSlotLeaderNums(ceilAvg);
        logger.info("[assign][assignLeaderSlots] begin, dataServers={}, avg={}, high={}, migrate={}", dataStoreNodeIdSet.size(), ceilAvg, highWaterMarkSlotLeaderNums, migratableSlotInfo);
        if (tryAssignLeaderSlots(highWaterMarkSlotLeaderNums)) {
            logger.info("[assign][after assignLeaderSlots] end -- leader changed");
            slotsBuilder.incrEpoch();
        } else {
            logger.info("[assign][after assignLeaderSlots] end -- no changes");
        }

        logger.info("[assign][assignFollowerSlots] begin");
        if (assignFollowerSlots()) {
            logger.info("[assign][after assignFollowerSlots] end -- follower changed");
            slotsBuilder.incrEpoch();
        } else {
            logger.info("[assign][assignFollowerSlots] end -- no changes");
        }
        return slotsBuilder.build();
    }

    private boolean tryAssignLeaderSlots(int highWaterMarkSlotLeaderNums) {
        /**
         * our strategy(assign leader) is to swap follower to leader when follower is enabled if no
         * followers, we select a new data-server to assign, that's simple and low prioritized so,
         * leaders are resolved in an order that who has least followers first (as we wish to satisfy
         * these nodes first) leaders with no follower is lowest priority, as whatever we did, it will
         * pick up a candidate that is not its follower
         */
        List<Integer> leaderSlotIdList = migratableSlotInfo.getLeadersByScore(new FewerFollowerFirstStrategy(slotsBuilder));
        if (leaderSlotIdList.isEmpty()) {
            logger.info("[assignLeaderSlots] no slot leader needs assign, quit");
            return false;
        }
        for (int leaderSlotId : leaderSlotIdList) {
            List<String> dataStoreNodeIdList = new ArrayList<>(dataStoreNodeIdSet);
            String nextLeaderDataStoreNodeId = Selectors.slotLeaderSelector(highWaterMarkSlotLeaderNums, slotsBuilder, leaderSlotId).select(dataStoreNodeIdList);
            boolean nextLeaderDataStoreNodeIdIsFollowerOfSlot = nextLeaderDataStoreNodeIdIsFollowerOfSlot(leaderSlotId, nextLeaderDataStoreNodeId);
            logger.info("[assignLeaderSlots]assign slot[{}] leader as [{}], upgrade={}, dataServers={}", leaderSlotId, nextLeaderDataStoreNodeId, nextLeaderDataStoreNodeIdIsFollowerOfSlot, dataStoreNodeIdSet.size());
            slotsBuilder.replaceLeaderDataStoreNodeId(leaderSlotId, nextLeaderDataStoreNodeId);
            Metrics.SlotAssign.onSlotLeaderAssign(nextLeaderDataStoreNodeId, leaderSlotId);
            if (nextLeaderDataStoreNodeIdIsFollowerOfSlot) {
                migratableSlotInfo.addLackFollowerNum(leaderSlotId);
            }
        }
        return true;
    }

    private boolean nextLeaderDataStoreNodeIdIsFollowerOfSlot(int slotId, String nextLeaderDataStoreNodeId) {
        return slotsBuilder.getOrCreate(slotId).containsFollowerDataStoreNodeId(nextLeaderDataStoreNodeId);
    }

    private boolean assignFollowerSlots() {
        List<FollowerToAssign> followerToAssignList = migratableSlotInfo.getFollowersByScore(new FollowerEmergentScoreJury());
        if (followerToAssignList.isEmpty()) {
            logger.info("[assignFollowerSlots] no follower slots need to assign, quit");
            return false;
        }
        int assignCount = 0;
        for (FollowerToAssign followerToAssign : followerToAssignList) {
            final int slotId = followerToAssign.getSlotId();
            for (int i = 0; i < followerToAssign.getAssigneeNums(); i++) {
                List<String> candidatesDataStoreNodeIdSList = new ArrayList<>(dataStoreNodeIdSet);
                candidatesDataStoreNodeIdSList.sort(Comparators.leastFollowersFirst(slotsBuilder));
                boolean assigned = false;
                for (String candidatesDataStoreNodeId : candidatesDataStoreNodeIdSList) {
                    DataStoreNode candidatesDataStoreNode = slotsBuilder.getDataStoreNode(candidatesDataStoreNodeId);
                    if (candidatesDataStoreNode.containsFollowerSlotId(slotId) || candidatesDataStoreNode.containsLeaderSlotId(slotId)) {
                        continue;
                    }
                    slotsBuilder.addFollowerDataStoreNodeId(slotId, candidatesDataStoreNodeId);
                    Metrics.SlotAssign.onSlotFollowerAssign(candidatesDataStoreNodeId, slotId);
                    assigned = true;
                    assignCount++;
                    logger.info("[assignFollowerSlots]assign slot[{}] add follower as [{}], dataServers={}", slotId, candidatesDataStoreNodeId, dataStoreNodeIdSet.size());
                    break;
                }
                if (!assigned) {
                    logger.warn("[assignFollowerSlots]assign slot[{}] no dataServer to assign, dataServers={}", slotId, dataStoreNodeIdSet.size());
                }
            }
        }
        return assignCount != 0;
    }

}
